#include "stdafx.h"
#include "SearchCode.h"
#include <stdio.h>

void CSearchCode::Search(CODE_SEARCH* pSearch, const char* logfilename)
{
    SYSTEM_INFO system_info;
    GetSystemInfo(&system_info);
    BYTE *pMaxAddress=(BYTE*)system_info.lpMaximumApplicationAddress;

    int anz_functions=0;
    { // clear all function pointers in pSearch
        CODE_SEARCH* pS=pSearch;
        while(!pS->IsTerminator())
        {
            *pS->ppFunction=NULL;
            if(pS->IsValid()) anz_functions++;
            pS++;
        }
    }

    union
    {
        BYTE* pCurrent;
        DWORD current_dw;
    };
    pCurrent=(BYTE*)system_info.lpMinimumApplicationAddress;

    int anz_found=0;
    int last_page=-1;

    FILE* pFile=fopen(logfilename,"w");
    fprintf(pFile," // Starting LUA search for %d functions at [%p - %p]\n",anz_functions,pCurrent,pMaxAddress);
    fprintf(pFile,"#include <idc.idc>\nstatic main()\n{\n");

    while((!pSearch->IsTerminator()) && (pCurrent<pMaxAddress))
    {
        if(pSearch->IsValid())
        {
            int code_len=0;
            BYTE code_search[MAX_CODE_LEN];
            BYTE code_ignore[MAX_CODE_LEN];

            const char* pCodeSearch=pSearch->pCodeSearch;
            const char* pCodeIgnore=pSearch->pCodeIgnore;

            for(code_len=0;code_len<MAX_CODE_LEN;code_len++)
            {
                if(!pCodeIgnore[code_len]) break;
                code_ignore[code_len]= pCodeIgnore[code_len]=='?' ? 0 : 0xff;
                code_search[code_len]= pCodeSearch[code_len]&code_ignore[code_len];
            }

            if(!code_len)
            {
                fprintf(pFile,"// No Code for %s\n",pSearch->pName);
            }
            else
            {
                fprintf(pFile,"// Searching for %s codelength=%d : ",pSearch->pName,code_len);
                for(int i=0;i<code_len;i++)
                {
                    if(code_ignore[i]) fprintf(pFile,"%02X ",code_search[i]);
                    else               fprintf(pFile,"?? ");
                }
                fprintf(pFile,"\n");

                BOOL found=FALSE;
                do
                {
                    int page=current_dw/system_info.dwPageSize;
                    if(page!=last_page)
                    {
                        last_page=page;
                        MEMORY_BASIC_INFORMATION mem_info;
                        VirtualQuery(pCurrent,&mem_info,sizeof(MEMORY_BASIC_INFORMATION));
                        while( (mem_info.State!=MEM_COMMIT)||(mem_info.Protect!=PAGE_EXECUTE_READ) )
                        {   // try next page
                            current_dw+=system_info.dwPageSize;
                            current_dw&=~(system_info.dwPageSize-1);
                            if(pCurrent>=pMaxAddress) break;
                            VirtualQuery(pCurrent,&mem_info,sizeof(MEMORY_BASIC_INFORMATION));
                        }
                    }

                    if(pCurrent>=pMaxAddress) break;

                    int compare;
                    for(compare=0;compare<code_len;compare++)
                    {
                        if((pCurrent[compare]&code_ignore[compare])!=code_search[compare]) 
                        {
                            break;
                        }
                    }

                    if(compare==code_len)
                    {
                        fprintf(pFile,"   MakeName(0x%p,\"%s\");\n",pCurrent,pSearch->pName);
                        *pSearch->ppFunction=(void*)pCurrent;
                        pCurrent+=code_len;
                        anz_found++;
                        found=TRUE;
                    }
                    else
                    {
                        pCurrent++;
                    }
                    // align the pointer
                    current_dw+=0xf;
                    current_dw&=~0xf;
                }while(!found);
            }
        }
        pSearch++;
    }
    fprintf(pFile,"}\n");
    fprintf(pFile,"\n// Ended LUA search , found %d\n",anz_found);
    if(anz_found==anz_functions) fprintf(pFile,"// Found all functions\n");
    else                         fprintf(pFile,"// Missing functions\n");
    fclose(pFile);
}